import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.db.session import SessionLocal, Base, engine
from app.models.graph import Node, Relation, Property, NodeType, RelationType, PropertyType
import uuid
from datetime import datetime

# 创建数据库表
Base.metadata.create_all(bind=engine)

def create_test_data():
    db = SessionLocal()
    try:
        # 创建人物节点
        person1 = Node(
            id=str(uuid.uuid4()),
            type=NodeType.PERSON,
            name="爱因斯坦",
            description="著名物理学家，相对论的创立者"
        )
        db.add(person1)

        person2 = Node(
            id=str(uuid.uuid4()),
            type=NodeType.PERSON,
            name="牛顿",
            description="著名物理学家，经典力学的奠基人"
        )
        db.add(person2)

        # 创建概念节点
        concept1 = Node(
            id=str(uuid.uuid4()),
            type=NodeType.CONCEPT,
            name="相对论",
            description="描述时空和引力的物理学理论"
        )
        db.add(concept1)

        concept2 = Node(
            id=str(uuid.uuid4()),
            type=NodeType.CONCEPT,
            name="经典力学",
            description="研究物体运动规律的物理学分支"
        )
        db.add(concept2)

        # 创建组织节点
        org1 = Node(
            id=str(uuid.uuid4()),
            type=NodeType.ORGANIZATION,
            name="普林斯顿大学",
            description="美国著名研究型大学"
        )
        db.add(org1)

        # 添加属性
        for person in [person1, person2]:
            db.add(Property(
                id=str(uuid.uuid4()),
                node_id=person.id,
                key="出生日期",
                value="1879-03-14" if person == person1 else "1643-01-04",
                type=PropertyType.DATE
            ))
            db.add(Property(
                id=str(uuid.uuid4()),
                node_id=person.id,
                key="国籍",
                value="德国" if person == person1 else "英国",
                type=PropertyType.TEXT
            ))

        # 添加关系
        db.add(Relation(
            id=str(uuid.uuid4()),
            type=RelationType.CREATED,
            source_id=person1.id,
            target_id=concept1.id,
            weight=1
        ))

        db.add(Relation(
            id=str(uuid.uuid4()),
            type=RelationType.CREATED,
            source_id=person2.id,
            target_id=concept2.id,
            weight=1
        ))

        db.add(Relation(
            id=str(uuid.uuid4()),
            type=RelationType.BELONGS_TO,
            source_id=person1.id,
            target_id=org1.id,
            weight=1
        ))

        db.commit()
        print("测试数据创建成功！")

    except Exception as e:
        print(f"创建测试数据失败: {e}")
        db.rollback()
    finally:
        db.close()

if __name__ == "__main__":
    create_test_data() 